#include "Segmentation_GPU.h"

__device__ float getIntensityDiff(float4* one, float4* two){
	float x = one->x - two->x;
	float y = one->y - two->y;
	float z = one->z - two->z;
	
	return sqrt(x * x + y * y + z * z);
}


__device__ void setLabel(float4* current, float4* prev, int* label, const float THRESHOLD){
	if(getIntensityDiff(current, prev) < THRESHOLD)
		current->w = current->w >= *label ? *label : current->w;			
		
	*label = current->w;
}



__global__ void segmentVertical(float4* in, int width, int height, const float THRESHOLD){
	int x = threadIdx.x + (blockIdx.x * blockDim.x);
	float4* current;
	float4* prev = &in[x];
	int label = prev->w;
	

	for(int y = 1; y < height; y++){
		current = &in[y * width + x];
		setLabel(current, prev, &label, THRESHOLD);
		prev = current;
	}__syncthreads();

	prev = &in[(height - 1) * width + x];
	label = prev->w;

	for(int y = height - 2; y >= 0; y--){
		current = &in[y * width + x];
		setLabel(current, prev, &label, THRESHOLD);
		prev = current;
	}
}


__global__ void segmentHorizontal(float4* in, int width, int height, const float THRESHOLD){
	int y = threadIdx.x + (blockIdx.x * blockDim.x);
	float4* current;
	float4* prev = &in[y * width];
	int label = prev->w;

	for(int x = 1; x < width; x++){
		current = &in[y * width + x];
		setLabel(current, prev, &label, THRESHOLD);
		prev = current;
	}__syncthreads();

	prev = &in[y * width + (width - 1)];
	label = prev->w;

	for(int x = width - 2; x >= 0; x--){
		current = &in[y * width + x];
		setLabel(current, prev, &label, THRESHOLD);
		prev = current;
	}
}

void Segmentation_GPU::runInitialSegmentation(){
	for(int i = 0; i < 30; i++){
		segmentVertical<<<width / 32, 32>>>(clusterMap_C, width, height, THRESHOLD);	
		segmentHorizontal<<<height / 30, 30>>>(clusterMap_C, width, height, THRESHOLD);
	}
}

__global__ void calcClusterSize(float4* in, int* clusterMap, int width, int height){
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;
	int access = x + y * width;

	atomicAdd(&clusterMap[(int)in[access].w], 1);
}


__global__ void dropClusters(float4* in, int* clusterMap, int width, int height){
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;
	int access = x + y * width;

	int cluster = in[access].w;
	if(clusterMap[cluster] < 50)
		in[access].w = -1;	
}

void Segmentation_GPU::postProcess(){
	thrust::fill(
		thrust::device_ptr<int>(clusterMap),
		thrust::device_ptr<int>(clusterMap) + width * height,
		0
	);

	calcClusterSize<<<dim3(width / 32, height / 24), dim3(32, 24)>>>
		(clusterMap_C, clusterMap, width, height);


	dropClusters<<<dim3(width / 32, height / 24), dim3(32, 24)>>>
		(clusterMap_C, clusterMap, width, height);


}